from typing import Dict, List, Union
from geesibling.core.types import Graph, Node
from ..StrategiesVector import StrategiesVector
from .node_handler import NodeHandler
from .registry import operator_registry
from ..shard.placement_types import (
    AbstractTensor,DeviceMesh,OperationData, OperationDataType, ShardingStrategy
)
from ..generator import EmbeddingStrategyGenerator, StrategyGenerator

__all__ = ['EmbeddingModuleHandler', 'EmbeddingFunctionHandler']




@operator_registry.register('aten.embedding.default')
class EmbeddingModuleHandler(NodeHandler):
    """
    A EmbeddingModuleHandler which deals with the sharding strategies for nn.Embedding module.
    """
    def __init__(self,node: Node,device_mesh: DeviceMesh,strategies_vector: StrategiesVector,pre_strategies_vectors:List[StrategiesVector]):
        super().__init__(node, device_mesh, strategies_vector)  # 调用父类的构造方法
        self.pre_strategies_vectors = pre_strategies_vectors

    def get_strategy_generator(self) -> List[StrategyGenerator]:
        op_data_mapping = self.get_operation_data_mapping()
        generators = []
        generators.append(EmbeddingStrategyGenerator(op_data_mapping, self.device_mesh))
        return generators

    def get_operation_data_mapping(self) -> Dict[str, OperationData]:
        # bias input weight
        input_abs_data = AbstractTensor(self.node.input_shape(1),self.node.input_type(1))
        weight_abs_data = AbstractTensor(self.node.input_shape(0),self.node.input_type(0))
        output_abs_data = AbstractTensor(self.node.output_shape(0),self.node.output_type(0))
        op_data_mapping = {
            "input":OperationData(name=str(self.node.inputs[1]), data=input_abs_data,type=OperationDataType.INPUT),
            "weight":OperationData(name=str(self.node.inputs[0]), data=weight_abs_data,type=OperationDataType.PARAM),
            "output":OperationData(name=str(self.node.name), data=output_abs_data,type=OperationDataType.OUTPUT)
        }
        return op_data_mapping
        return mapping

